# -*- coding: utf-8 -*-
"""
@Time ： 2025/5/26 15:45
@Auth ： 李天天
"""
# !/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import pandas as pd
import numpy as np
import openpyxl
from openpyxl.styles import PatternFill, Font, Alignment
from openpyxl.utils import get_column_letter

# 定义颜色常量
GREEN = '\033[92m'  # 绿色
RED = '\033[91m'  # 红色
RESET = '\033[0m'  # 重置颜色


def calculate_error(input_file, excel_output=None):
    """计算误差值并输出结果到控制台和Excel(可选)"""
    # 读取输入文件
    df = pd.read_csv(input_file, sep='\t', header=None)
    df.columns = ['帧名称', '算法预测遮挡面积', '实际遮挡面积']

    # 计算误差值
    df['误差值'] = abs(df['算法预测遮挡面积'] - df['实际遮挡面积']) * 100

    # 设置档位
    thresholds = [5, 10, 15, 20]

    # 判断是否符合各个档位
    for threshold in thresholds:
        df[f'{threshold}%'] = df['误差值'].apply(lambda x: f'{GREEN}√{RESET}' if x <= threshold else f'{RED}x{RESET}')
        # 为Excel输出保存原始标记
        df[f'{threshold}%_raw'] = df['误差值'].apply(lambda x: True if x <= threshold else False)

    # 判断遮挡比例档位预测准确性
    def get_range(value):
        if 0 <= value <= 0.3:
            return '0-30%'
        elif 0.3 < value <= 0.7:
            return '31-70%'
        else:
            return '71-100%'

    df['预测档位'] = df['算法预测遮挡面积'].apply(get_range)
    df['实际档位'] = df['实际遮挡面积'].apply(get_range)
    df['档位预测'] = df.apply(
        lambda row: f'{GREEN}√{RESET}' if row['预测档位'] == row['实际档位'] else f'{RED}x{RESET}', axis=1)
    # 为Excel输出保存原始标记
    df['档位预测_raw'] = df.apply(lambda row: True if row['预测档位'] == row['实际档位'] else False, axis=1)

    # 计算各个档位的准确率
    total_rows = len(df)
    accuracy_stats = {}
    for threshold in thresholds:
        correct_count = sum(df['误差值'] <= threshold)
        accuracy = (correct_count / total_rows) * 100
        accuracy_stats[threshold] = accuracy

    # 计算遮挡比例档位预测准确率
    range_stats = {}
    for range_name in ['0-30%', '31-70%', '71-100%']:
        range_df = df[df['实际档位'] == range_name]
        total = len(range_df)
        correct = sum(range_df['预测档位'] == range_df['实际档位'])
        if total > 0:
            accuracy = (correct / total) * 100
        else:
            accuracy = 0
        range_stats[range_name] = {
            '总帧数': total,
            '正确预测帧数': correct,
            '准确率': accuracy
        }

    # 输出结果到控制台
    # 输出表头
    print('帧名称\t预测值\t实际值\t误差值\t5%\t10%\t15%\t20%\t档位预测')
    print('-' * 80)

    # 输出每一行数据
    for _, row in df.iterrows():
        print(f"{row['帧名称']}\t{row['算法预测遮挡面积']:.6f}\t{row['实际遮挡面积']:.6f}\t{row['误差值']:.2f}%\t",
              end='')
        for threshold in thresholds:
            print(f"{row[f'{threshold}%']}\t", end='')
        print(f"{row['档位预测']}")

    # 输出统计结果
    print('\n\n准确率统计：')
    for threshold in thresholds:
        print(f"误差值 ≤ {threshold}% 的准确率: {accuracy_stats[threshold]:.2f}%")

    # 输出遮挡比例档位预测准确率
    print('\n遮挡比例档位  总帧数 正确预测帧数    准确率')
    for range_name, stats in range_stats.items():
        print(f"{range_name}\t{stats['总帧数']}\t{stats['正确预测帧数']}\t{stats['准确率']:.2f}%")

    # 如果指定了Excel输出，则输出到Excel
    if excel_output:
        # 创建一个新的Excel工作簿和工作表
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = "误差分析结果"

        # 设置列宽
        for col in range(1, 10):
            ws.column_dimensions[get_column_letter(col)].width = 15

        # 设置单元格样式
        title_font = Font(bold=True)
        center_alignment = Alignment(horizontal='center', vertical='center')
        green_fill = PatternFill(start_color="92D050", end_color="92D050", fill_type="solid")
        red_fill = PatternFill(start_color="FF0000", end_color="FF0000", fill_type="solid")

        # 添加准确率统计到顶部
        row_idx = 1
        ws.cell(row=row_idx, column=1, value="准确率统计").font = title_font
        row_idx += 1

        for threshold in thresholds:
            ws.cell(row=row_idx, column=1, value=f"误差值 ≤ {threshold}% 的准确率:")
            ws.cell(row=row_idx, column=2, value=f"{accuracy_stats[threshold]:.2f}%")
            row_idx += 1

        row_idx += 1
        ws.cell(row=row_idx, column=1, value="遮挡比例档位预测准确率").font = title_font
        row_idx += 1

        # 添加遮挡比例档位预测准确率表头
        ws.cell(row=row_idx, column=1, value="遮挡比例档位").font = title_font
        ws.cell(row=row_idx, column=2, value="总帧数").font = title_font
        ws.cell(row=row_idx, column=3, value="正确预测帧数").font = title_font
        ws.cell(row=row_idx, column=4, value="准确率").font = title_font
        row_idx += 1

        # 添加遮挡比例档位预测准确率数据
        for range_name, stats in range_stats.items():
            ws.cell(row=row_idx, column=1, value=range_name)
            ws.cell(row=row_idx, column=2, value=stats['总帧数'])
            ws.cell(row=row_idx, column=3, value=stats['正确预测帧数'])
            ws.cell(row=row_idx, column=4, value=f"{stats['准确率']:.2f}%")
            row_idx += 1

        # 添加空行
        row_idx += 2

        # 添加数据表头
        headers = ['帧名称', '预测值', '实际值', '误差值', '5%', '10%', '15%', '20%', '档位预测']
        for col_idx, header in enumerate(headers, 1):
            ws.cell(row=row_idx, column=col_idx, value=header).font = title_font
            ws.cell(row=row_idx, column=col_idx).alignment = center_alignment
        row_idx += 1

        # 添加数据行
        for _, row in df.iterrows():
            ws.cell(row=row_idx, column=1, value=row['帧名称'])
            ws.cell(row=row_idx, column=2, value=row['算法预测遮挡面积'])
            ws.cell(row=row_idx, column=3, value=row['实际遮挡面积'])
            ws.cell(row=row_idx, column=4, value=f"{row['误差值']:.2f}%")

            # 添加档位标记，并设置颜色
            for col_idx, threshold in enumerate(thresholds, 5):
                ws.cell(row=row_idx, column=col_idx, value="√" if row[f'{threshold}%_raw'] else "×")
                ws.cell(row=row_idx, column=col_idx).alignment = center_alignment
                if row[f'{threshold}%_raw']:
                    ws.cell(row=row_idx, column=col_idx).fill = green_fill
                else:
                    ws.cell(row=row_idx, column=col_idx).fill = red_fill

            # 添加档位预测标记，并设置颜色
            ws.cell(row=row_idx, column=9, value="√" if row['档位预测_raw'] else "×")
            ws.cell(row=row_idx, column=9).alignment = center_alignment
            if row['档位预测_raw']:
                ws.cell(row=row_idx, column=9).fill = green_fill
            else:
                ws.cell(row=row_idx, column=9).fill = red_fill

            row_idx += 1

        # 保存Excel文件
        wb.save(excel_output)
        print(f"\nExcel文件已保存到: {excel_output}")


def main():
    if len(sys.argv) > 1:
        input_file = sys.argv[1]
    else:
        # 默认使用merged_output2.txt作为输入
        input_file = 'merged_output2.txt'

    # 生成与输入文件同名但扩展名为.xlsx的Excel输出文件名
    base_name = os.path.splitext(os.path.basename(input_file))[0]
    excel_output = f"{base_name}.xlsx"

    calculate_error(input_file, excel_output)
    print("\n计算完成")


if __name__ == "__main__":
    main()